Hey @tengyifei! Let me know if you're able to review this, or feel free to add anyone else. We have this as a requirement to unblock our use case above, and we might consider cherry picking it to 2.7.1, since it is a blocker.

@bhavya01 @zhanyong-wan Do you have cycles on this? We need this to unblock some of our use cases.

@bhavya01 @ysiraichi Hey folks, stale PR - let me know if you can review.

I will take a look at this later today!

I am not really familiar with sharding, but I do have one question: wouldn't it be better to allow mark_sharding to be called multiple times, instead of creating a new API? I'm asking this question because, looking at the test, it's almost like mark_sharding and annotate_custom_sharding did the same thing (i.e. would a user know when to use which?).

At the moment, mark_sharding very simply does; 1/ If the tensor is an IR value (and not a clone of a device data node), add an IR custom sharding node, or 2/ If the tensor is a Device Data node on the host/device, we will actually invoke the sharded data transfer to the device.

This means that it's not possible to provide a custom sharding annotation (on XLA) over an intentionally replicated tensor on the device. Similarly, as the motivation above mentions, it also limits us to provide a custom sharding annotation over a Device Data node (weights, inputs, etc) that have already been sharded to the device - since it can't disambiguate if a user meant to reshard the tensor, or simply add a custom sharding op.

This API intended to provide these gaps to a user - and is intentionally targeting more familiar users who need to provide extra sharding annotation to XLA around the limitations above on their tensors. I think it's a well defined API that can be served for different use cases. If you have a new tensor on the host, mark_sharding and annotate_custom_sharding will have different reasonable behaviors.

I am happy to consider relaxing the mark_sharding API, e.g. removing: https://github.com/pytorch/xla/pull/9203/files#diff-76bd84e4abe22701ee8697bf77e9fc97e19b6d6ff05175f2dc87f938f3a88837R774-R775, but that basically means that if a tensor has already been sharded on the device, any following mark_sharding will entail an XLA custom sharding annotation - and NOT an attempt to reshard the tensor (which fails today if the sharding spec differs). Perhaps we introduce an argument that makes it clearer?

@ysiraichi Do you want to take another look before merging?

@ysiraichi Let me know if the above is justified.

Sorry for taking too long.

Perhaps we introduce an argument that makes it clearer?

That's what I was thinking. That said, again, I'm not familiar with sharding, so I trust your judgement, here. Also, thank you for the thorough explanation. It was very clarifying.

Sorry for taking too long.

Perhaps we introduce an argument that makes it clearer?

That's what I was thinking. That said, again, I'm not familiar with sharding, so I trust your judgement, here. Also, thank you for the thorough explanation. It was very clarifying.

No worries, thanks for taking the time. I think, for now, it's comparably better to introduce the API given the different behavior, than to modify/extend an essential standard one. I am happy to eventually converge as we see fit.

I'll rebase and re-run the CI, thanks folks!